import os.path
import time

import cv2


def save_rtsp_video(rtsp_url, base_path):
    cap = cv2.VideoCapture(rtsp_url)
    if not cap.isOpened():
        print(f"无法打开 RTSP 流：{rtsp_url}")
        return
    fourcc = cv2.VideoWriter.fourcc('m', 'p', '4', 'v')
    frame_s = cap.get(cv2.CAP_PROP_FPS)
    size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    time_frame = frame_s * 60 * 60  # 设置保存时间为1小时
    frame_num = 0
    video_writer = None
    while True:
        # 参数：视频文件名，格式，每秒帧数，宽高，是否灰度
        ret, frame = cap.read()
        if not ret:
            print(f"无法读取帧，尝试重新连接 RTSP 流：{rtsp_url}")
            cap.release()
            time.sleep(5)
            cap = cv2.VideoCapture(rtsp_url)
            if not cap.isOpened():
                continue
        if frame_num == 0:
            filename = append_file_name(base_path)
            video_writer = cv2.VideoWriter(filename, fourcc, frame_s, size, True)
        video_writer.write(frame)
        frame_num += 1
        if frame_num == time_frame:
            video_writer.release()
            frame_num = 0
    cap.release()


def append_file_name(base_path):
    now = time.localtime()
    year = now.tm_year
    month = now.tm_mon
    day = now.tm_mday
    hour = now.tm_hour
    dir_path = f"{base_path}/{year}/{month}/{day}"
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)
    return f'{dir_path}/{hour}.mp4'


if __name__ == '__main__':
    save_rtsp_video("rtsp://chi-square.com.cn:6004/live/test", "d:/ffmpeg")
